"""
Comprehensive boundary condition handling for all differential operators.

This module provides a unified interface for applying various boundary conditions
to differential operators across all coordinate systems. It supports:

- Dirichlet (fixed value): u = g
- Neumann (fixed derivative): ∂u/∂n = g  
- Robin (mixed): au + b∂u/∂n = g
- Periodic: u(x₀) = u(x₁), ∂u/∂n(x₀) = ∂u/∂n(x₁)
- Absorbing/Radiation boundary conditions
- Custom user-defined boundary conditions

Features:
- Automatic stencil modification for different BC types
- Ghost point handling for high-order schemes  
- Time-dependent boundary conditions
- Vector and tensor field boundary conditions
- Compatibility with all coordinate systems
"""

using LinearAlgebra
using SparseArrays

# Re-export boundary condition types from common.jl
export DirichletBC, NeumannBC, RobinBC, PeriodicBC
export AbsorbingBC, RadiationBC, CustomBC
export BoundaryConditionSet, apply_boundary_conditions!
export modify_operator_for_bc!, create_ghost_points!
export time_dependent_bc, vector_field_bc

"""
    AbsorbingBC

Absorbing boundary condition for wave equations: ∂u/∂t + c∂u/∂n = 0
"""
struct AbsorbingBC <: BoundaryCondition
    wave_speed::Float64
end

"""
    RadiationBC

Radiation boundary condition: ∂u/∂r - (iku + α/r)u = 0
"""
struct RadiationBC <: BoundaryCondition
    wave_number::Complex{Float64}
    alpha::Float64
end

"""
    CustomBC

User-defined boundary condition with custom stencil modification function.
"""
struct CustomBC <: BoundaryCondition
    apply_function::Function
    description::String
end

"""
    BoundaryConditionSet

Container for boundary conditions on all faces/edges of a domain.
"""
struct BoundaryConditionSet
    conditions::Vector{BoundaryCondition}
    grid_type::Type
    dimension_names::Vector{Symbol}
    
    function BoundaryConditionSet(conditions::Vector{BoundaryCondition}, 
                                grid_type::Type, 
                                dimension_names::Vector{Symbol})
        new(conditions, grid_type, dimension_names)
    end
end

"""
    apply_boundary_conditions!(matrix::SparseMatrixCSC, rhs_vector::Vector, 
                              bc_set::BoundaryConditionSet, grid)

Apply all boundary conditions in the set to the operator matrix and RHS vector.
"""
function apply_boundary_conditions!(matrix::SparseMatrixCSC, rhs_vector::Vector, 
                                  bc_set::BoundaryConditionSet, grid)
    grid_type = typeof(grid)
    
    if grid_type <: CartesianGrid
        apply_cartesian_bc!(matrix, rhs_vector, bc_set, grid)
    elseif grid_type <: SphericalSymGrid
        apply_spherical_bc!(matrix, rhs_vector, bc_set, grid)
    elseif grid_type <: CylindricalSymGrid
        apply_cylindrical_bc!(matrix, rhs_vector, bc_set, grid)
    elseif grid_type <: PolarSymGrid
        apply_polar_bc!(matrix, rhs_vector, bc_set, grid)
    else
        error("Unsupported grid type for boundary conditions: $grid_type")
    end
    
    return matrix, rhs_vector
end

"""
    apply_cartesian_bc!(matrix, rhs_vector, bc_set, grid)

Apply boundary conditions for Cartesian grids.
"""
function apply_cartesian_bc!(matrix, rhs_vector, bc_set, grid)
    dims = size(grid)
    spacing_vals = spacing(grid)
    ndims = ndim(grid)
    
    # Linear index helper
    function linear_index(indices::Tuple)
        return LinearIndices(dims)[CartesianIndex(indices...)]
    end
    
    # Apply boundary conditions for each dimension
    for dim in 1:ndims
        bc_low = bc_set.conditions[2*dim - 1]  # Lower boundary
        bc_high = bc_set.conditions[2*dim]     # Upper boundary
        h = spacing_vals[dim]
        
        # Lower boundary (e.g., x=0, y=0, z=0)
        apply_bc_to_face!(matrix, rhs_vector, bc_low, grid, dim, :lower, h, linear_index)
        
        # Upper boundary (e.g., x=L, y=L, z=L)  
        apply_bc_to_face!(matrix, rhs_vector, bc_high, grid, dim, :upper, h, linear_index)
    end
end

"""
    apply_bc_to_face!(matrix, rhs_vector, bc, grid, dim, face, h, index_func)

Apply boundary condition to a specific face of the domain.
"""
function apply_bc_to_face!(matrix, rhs_vector, bc, grid, dim, face, h, index_func)
    dims = size(grid)
    ndims = ndim(grid)
    
    # Get indices for the boundary face
    boundary_indices = get_boundary_indices(dims, dim, face)
    
    for boundary_idx in boundary_indices
        i = index_func(boundary_idx)
        
        if isa(bc, DirichletBC)
            apply_dirichlet_bc!(matrix, rhs_vector, i, bc.value)
            
        elseif isa(bc, NeumannBC)
            apply_neumann_bc!(matrix, rhs_vector, i, bc.value, h, dim, face, 
                            boundary_idx, dims, index_func)
            
        elseif isa(bc, RobinBC)
            apply_robin_bc!(matrix, rhs_vector, i, bc.a, bc.b, bc.value, h, 
                          dim, face, boundary_idx, dims, index_func)
            
        elseif isa(bc, PeriodicBC)
            apply_periodic_bc!(matrix, rhs_vector, i, dim, boundary_idx, dims, index_func)
            
        elseif isa(bc, CustomBC)
            bc.apply_function(matrix, rhs_vector, i, grid, dim, face, boundary_idx)
        end
    end
end

"""
    get_boundary_indices(dims, dim, face)

Get all indices on a boundary face.
"""
function get_boundary_indices(dims, dim, face)
    ndims = length(dims)
    indices = Vector{Tuple}()
    
    # Create iterator over all boundary points on the specified face
    ranges = [1:dims[d] for d in 1:ndims]
    
    if face == :lower
        ranges[dim] = 1:1  # First slice in dimension dim
    else  # face == :upper
        ranges[dim] = dims[dim]:dims[dim]  # Last slice in dimension dim
    end
    
    for idx in Iterators.product(ranges...)
        push!(indices, idx)
    end
    
    return indices
end

"""
    apply_dirichlet_bc!(matrix, rhs_vector, i, value)

Apply Dirichlet boundary condition: u_i = value
"""
function apply_dirichlet_bc!(matrix, rhs_vector, i, value)
    # Clear row and set to identity
    matrix[i, :] .= 0.0
    matrix[i, i] = 1.0
    rhs_vector[i] = value
end

"""
    apply_neumann_bc!(matrix, rhs_vector, i, value, h, dim, face, boundary_idx, dims, index_func)

Apply Neumann boundary condition: ∂u/∂n = value
"""
function apply_neumann_bc!(matrix, rhs_vector, i, value, h, dim, face, 
                         boundary_idx, dims, index_func)
    # Clear row
    matrix[i, :] .= 0.0
    
    # Set finite difference stencil for derivative
    if face == :lower
        # Forward difference: (u_{i+1} - u_i)/h = value
        matrix[i, i] = -1.0/h
        
        # Get neighbor index
        neighbor_idx = collect(boundary_idx)
        neighbor_idx[dim] += 1
        j = index_func(tuple(neighbor_idx...))
        matrix[i, j] = 1.0/h
        
    else  # face == :upper
        # Backward difference: (u_i - u_{i-1})/h = value  
        matrix[i, i] = 1.0/h
        
        # Get neighbor index
        neighbor_idx = collect(boundary_idx)
        neighbor_idx[dim] -= 1
        j = index_func(tuple(neighbor_idx...))
        matrix[i, j] = -1.0/h
    end
    
    rhs_vector[i] = value
end

"""
    apply_robin_bc!(matrix, rhs_vector, i, a, b, value, h, dim, face, boundary_idx, dims, index_func)

Apply Robin boundary condition: au + b∂u/∂n = value
"""
function apply_robin_bc!(matrix, rhs_vector, i, a, b, value, h, dim, face, 
                       boundary_idx, dims, index_func)
    # Clear row
    matrix[i, :] .= 0.0
    
    # Robin condition: au_i + b(∂u/∂n)_i = value
    if face == :lower
        # Forward difference for derivative
        matrix[i, i] = a - b/h
        
        neighbor_idx = collect(boundary_idx)
        neighbor_idx[dim] += 1
        j = index_func(tuple(neighbor_idx...))
        matrix[i, j] = b/h
        
    else  # face == :upper
        # Backward difference for derivative
        matrix[i, i] = a + b/h
        
        neighbor_idx = collect(boundary_idx)
        neighbor_idx[dim] -= 1
        j = index_func(tuple(neighbor_idx...))
        matrix[i, j] = -b/h
    end
    
    rhs_vector[i] = value
end

"""
    apply_periodic_bc!(matrix, rhs_vector, i, dim, boundary_idx, dims, index_func)

Apply periodic boundary condition: u(x_0) = u(x_L)
"""
function apply_periodic_bc!(matrix, rhs_vector, i, dim, boundary_idx, dims, index_func)
    # Clear row
    matrix[i, :] .= 0.0
    
    # Connect boundary points: u_i - u_j = 0
    matrix[i, i] = 1.0
    
    # Find corresponding point on opposite boundary
    opposite_idx = collect(boundary_idx)
    if boundary_idx[dim] == 1
        opposite_idx[dim] = dims[dim]
    else
        opposite_idx[dim] = 1
    end
    
    j = index_func(tuple(opposite_idx...))
    matrix[i, j] = -1.0
    
    rhs_vector[i] = 0.0
end

"""
    modify_operator_for_bc!(operator_func, bc_set, grid)

Modify an existing operator function to properly handle boundary conditions.
"""
function modify_operator_for_bc!(operator_func, bc_set, grid)
    function bc_modified_operator(input, output)
        # Apply base operator
        operator_func(input, output)
        
        # Apply boundary condition corrections
        apply_bc_corrections!(output, input, bc_set, grid)
        
        return output
    end
end

"""
    apply_bc_corrections!(output, input, bc_set, grid)

Apply boundary condition corrections to operator output.
"""
function apply_bc_corrections!(output, input, bc_set, grid)
    dims = size(grid)
    ndims = ndim(grid)
    
    # Apply corrections for each boundary
    for dim in 1:ndims
        bc_low = bc_set.conditions[2*dim - 1]
        bc_high = bc_set.conditions[2*dim]
        
        apply_bc_correction_to_face!(output, input, bc_low, grid, dim, :lower)
        apply_bc_correction_to_face!(output, input, bc_high, grid, dim, :upper)
    end
end

"""
    apply_bc_correction_to_face!(output, input, bc, grid, dim, face)

Apply boundary condition correction to a specific face.
"""
function apply_bc_correction_to_face!(output, input, bc, grid, dim, face)
    dims = size(grid)
    boundary_indices = get_boundary_indices(dims, dim, face)
    
    for boundary_idx in boundary_indices
        idx = CartesianIndex(boundary_idx...)
        
        if isa(bc, DirichletBC)
            # Set output to specified value
            output[idx] = bc.value
            
        elseif isa(bc, NeumannBC)
            # Modify output to satisfy derivative condition
            # (Implementation depends on specific operator)
            continue
            
        elseif isa(bc, PeriodicBC)
            # Ensure periodicity
            if face == :lower
                opposite_idx = collect(boundary_idx)
                opposite_idx[dim] = dims[dim]
                output[idx] = input[CartesianIndex(opposite_idx...)]
            end
        end
    end
end

"""
    create_ghost_points!(field_data, bc_set, grid, ghost_width=1)

Create ghost points for high-order finite difference schemes.
"""
function create_ghost_points!(field_data, bc_set, grid, ghost_width=1)
    dims = size(grid)
    ndims = ndim(grid)
    
    # Extend field_data with ghost points
    extended_dims = tuple([dims[i] + 2*ghost_width for i in 1:ndims]...)
    extended_data = zeros(extended_dims)
    
    # Copy interior data
    interior_ranges = tuple([ghost_width+1:dims[i]+ghost_width for i in 1:ndims]...)
    extended_data[interior_ranges...] = field_data
    
    # Fill ghost points according to boundary conditions
    for dim in 1:ndims
        bc_low = bc_set.conditions[2*dim - 1]
        bc_high = bc_set.conditions[2*dim]
        
        fill_ghost_points!(extended_data, bc_low, grid, dim, :lower, ghost_width)
        fill_ghost_points!(extended_data, bc_high, grid, dim, :upper, ghost_width)
    end
    
    return extended_data
end

"""
    fill_ghost_points!(extended_data, bc, grid, dim, face, ghost_width)

Fill ghost points for a specific boundary face.
"""
function fill_ghost_points!(extended_data, bc, grid, dim, face, ghost_width)
    extended_dims = size(extended_data)
    
    if isa(bc, DirichletBC)
        # Extrapolate to satisfy Dirichlet condition
        fill_dirichlet_ghosts!(extended_data, bc.value, dim, face, ghost_width)
        
    elseif isa(bc, NeumannBC)
        # Extrapolate to satisfy Neumann condition
        fill_neumann_ghosts!(extended_data, bc.value, grid, dim, face, ghost_width)
        
    elseif isa(bc, PeriodicBC)
        # Copy from opposite boundary
        fill_periodic_ghosts!(extended_data, dim, face, ghost_width)
    end
end

"""
    time_dependent_bc(bc_func::Function, t::Float64) -> BoundaryCondition

Create time-dependent boundary condition.
"""
function time_dependent_bc(bc_func::Function, t::Float64)
    value = bc_func(t)
    return DirichletBC(value)  # Could be generalized for other BC types
end

"""
    vector_field_bc(bc_conditions::Vector{BoundaryCondition}, component::Int) -> BoundaryCondition

Extract boundary condition for specific component of vector field.
"""
function vector_field_bc(bc_conditions::Vector{BoundaryCondition}, component::Int)
    if length(bc_conditions) >= component
        return bc_conditions[component]
    else
        # Default to homogeneous Dirichlet if not specified
        return DirichletBC(0.0)
    end
end

"""
    create_bc_set_cartesian_2d(bc_left, bc_right, bc_bottom, bc_top) -> BoundaryConditionSet

Convenience function to create boundary condition set for 2D Cartesian grid.
"""
function create_bc_set_cartesian_2d(bc_left, bc_right, bc_bottom, bc_top)
    conditions = [bc_left, bc_right, bc_bottom, bc_top]
    return BoundaryConditionSet(conditions, CartesianGrid, [:x, :y])
end

"""
    create_bc_set_spherical_1d(bc_center, bc_boundary) -> BoundaryConditionSet

Convenience function for spherically symmetric problems.
"""
function create_bc_set_spherical_1d(bc_center, bc_boundary)
    conditions = [bc_center, bc_boundary]
    return BoundaryConditionSet(conditions, SphericalSymGrid, [:r])
end

"""
    validate_bc_compatibility(bc_set::BoundaryConditionSet, operator_type::Symbol)

Validate that boundary conditions are compatible with the operator type.
"""
function validate_bc_compatibility(bc_set::BoundaryConditionSet, operator_type::Symbol)
    for (i, bc) in enumerate(bc_set.conditions)
        if operator_type == :laplacian
            # All standard BCs are compatible with Laplacian
            if !isa(bc, Union{DirichletBC, NeumannBC, RobinBC, PeriodicBC})
                @warn "Boundary condition at position $i may not be compatible with Laplacian operator"
            end
        elseif operator_type == :gradient
            # Periodic BCs require special handling for gradient
            if isa(bc, PeriodicBC)
                @warn "Periodic boundary conditions require special handling for gradient operators"
            end
        end
    end
    
    return true
end

# Helper functions for specific ghost point filling
function fill_dirichlet_ghosts!(extended_data, value, dim, face, ghost_width)
    # Implementation for Dirichlet ghost point extrapolation
    # This would be specific to the grid layout
end

function fill_neumann_ghosts!(extended_data, derivative_value, grid, dim, face, ghost_width)
    # Implementation for Neumann ghost point extrapolation
    h = spacing(grid)[dim]
    # Use finite difference to satisfy derivative condition
end

function fill_periodic_ghosts!(extended_data, dim, face, ghost_width)
    # Implementation for periodic ghost points
    # Copy from opposite end of domain
end